# Neural Net Definition, Standard CNNs
import torch.nn as nn
from torch.nn import init


class SiameseNetwork(nn.Module):
    def __init__(self):
        super(SiameseNetwork, self).__init__()
        # self.cnn1 = nn.Sequential(
        #     nn.ReflectionPad2d(1),
        #     nn.Conv2d(1, 4, kernel_size=3),
        #     nn.ReLU(inplace=True),
        #     nn.BatchNorm2d(4),
        #     #nn.Dropout2d(p=.2),
        #
        #     nn.ReflectionPad2d(1),
        #     nn.Conv2d(4, 8, kernel_size=3),
        #     nn.ReLU(inplace=True),
        #     nn.BatchNorm2d(8),
        #     #nn.Dropout2d(p=.2),
        #
        #     # nn.ReflectionPad2d(1),
        #     # nn.Conv2d(8, 8, kernel_size=3),
        #     # nn.ReLU(inplace=True),
        #     # nn.BatchNorm2d(8),
        #     # nn.Dropout2d(p=.2),
        # )
        #
        # self.fc1 = nn.Sequential(
        #     nn.Linear(2 * 100 * 100, 2000),
        #     nn.ReLU(inplace=True),
        #
        #     nn.Linear(2000, 500),
        #     nn.ReLU(inplace=True),
        #
        #     nn.Linear(500, 5)
        # )

        self.cnn1 = nn.Sequential(
            nn.Conv2d(3, 50, kernel_size=7),
            nn.MaxPool2d(2, stride=2),
            nn.Conv2d(50, 100, kernel_size=7),
            nn.MaxPool2d(2, stride=2))

        self.fc1 = nn.Sequential(
            nn.Linear(16900, 1690),
            nn.ReLU(inplace=True),
            nn.Linear(1690, 40),
            nn.ReLU(inplace=True),
            nn.Linear(40, 2))
        self._initialize_weights()

    def _initialize_weights(self):
        print('_initialize_weights')
        # print(self.modules())

        for m in self.modules():
            # print(m)
            if isinstance(m, nn.Linear):
                # print(m.weight.data.type())
                # input()
                # m.weight.data.fill_(1.0)
                init.xavier_uniform_(m.weight, gain=1)

    def forward_once(self, x):
        output = self.cnn1(x)
        output = output.view(output.size()[0], -1)
        output = self.fc1(output)
        return output

    def forward(self, input1, input2):
        output1 = self.forward_once(input1)
        output2 = self.forward_once(input2)
        return output1, output2
